'''
利用均线排列情况，每周定投一次
引入一个额外的状态变量来跟踪上次买入的日期，确保在新的一周开始时才执行买入操作。
'''

from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
#stock_hfq_df = ak.stock_zh_a_hist(symbol="000001", adjust="hfq").iloc[:, :6]  #深圳股票
stock_hfq_df= ak.fund_etf_hist_em(symbol="510500", adjust="").iloc[:, :6]  #etf


# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])

#==========================================================================================
class MyStrategy(bt.Strategy):
    def __init__(self):
        self.sma5 = bt.indicators.MovingAverageSimple(self.data.close, period=5)
        self.sma30 = bt.indicators.MovingAverageSimple(self.data.close, period=30)
        self.last_buy_date = None  # 新增状态变量，记录上次买入的日期
        self.sum_buy=0 #累计买入资金
        self.buy_days=7 #间隔几天买入
        

    def next(self):
        current_date = self.data.datetime.date(0)
        # 检查是否是新的一周（假设每周一执行买入）
        if current_date.weekday() == 0 and (self.last_buy_date is None or (current_date - self.last_buy_date).days >= self.buy_days):
            if self.sma5 < self.sma30:  # 当SMA5小于SMA30时
                # 计算可购买的股数，这里简单地设定为100股
                size = 1000
                cash_available = self.broker.getcash()  # 获取可用资金
                price = self.data.close[0]  # 当前价格
                # 确保有足够的资金购买至少一股
                if cash_available >= price * size:
                    self.buy(size=size)  # 执行买入100股的操作
                    self.last_buy_date = current_date  # 更新上次买入的日期
                    #累计投资金额
                    self.sum_buy+=price * size #/100是换算成元
                    print(f"{current_date}: 买入价格: {price:.2f},  累计买入资金{self.sum_buy:.2f}")
                else:
                    print(f"{current_date}: 余额不足，无法购买。")
                

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))
#=======================================================================================================

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(2023, 6, 3)
end_date = datetime(2024, 5, 27)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MyStrategy)

# 设置初始资金和手续费
start_cash = 100000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")


# 绘制图表
cerebro.plot(style='candlestick',figsize=(12, 8))
